import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import matplotlib.dates as mdates
from matplotlib.colors import LinearSegmentedColormap

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


class CropIrrigationPlanner:
    """问题4：作物月度灌溉方案规划"""

    def __init__(self, problem2_result, problem3_result):
        # 1. 作物参数（表1数据）
        self.crops = {
            '高粱': {
                'area': 0.5,  # 种植面积（公顷）=5000平方米
                'sowing_days': 20,  # 播种期天数
                'flowering_days': 50,  # 开花期天数
                'mature_days': 20,  # 成熟期天数
                'sowing_water': 5,  # 播种期需水量（L/m²/天）
                'flowering_water': 10,  # 开花期需水量（L/m²/天）
                'mature_water': 8  # 成熟期需水量（L/m²/天）
            },
            '玉米': {
                'area': 0.3,  # 3000平方米
                'sowing_days': 32,
                'flowering_days': 50,
                'mature_days': 20,
                'sowing_water': 6,
                'flowering_water': 12,
                'mature_water': 10
            },
            '大豆': {
                'area': 0.2,  # 2000平方米
                'sowing_days': 40,
                'flowering_days': 40,
                'mature_days': 20,
                'sowing_water': 4,
                'flowering_water': 8,
                'mature_water': 6
            }
        }

        # 2. 种植与生长时间线（播种日：2021-05-01）
        self.plant_date = datetime(2021, 5, 1)
        self.calculate_growth_periods()  # 计算各作物生长期

        # 3. 现有系统配置（问题2结果）
        self.system_config = {
            'pipe_daily_flow': problem2_result['管道日流量'],  # 管道日流量（L）
            'tank_volume': problem2_result['储水罐容积'],  # 储水罐总容积（L）
            'emergency_ratio': problem3_result['emergency_ratio'],  # 应急储备比例（问题3结果）
            'pipe_length': problem2_result['管道长度'],  # 管道长度（米）
            'sprinkler_coverage': problem2_result['喷头覆盖半径']  # 喷头覆盖半径（米）
        }

        # 4. 生成模拟数据（实际应从附件导入）
        self.soil_data = self.generate_soil_data()  # 土壤湿度数据
        self.rainfall_data = self.generate_rainfall_data()  # 降水量数据

        # 5. 时间范围（5月1日至7月31日）
        self.date_range = pd.date_range(start='2021-05-01', end='2021-07-31')

        # 6. 月度结果存储
        self.monthly_results = []

    def calculate_growth_periods(self):
        """计算各作物的生长周期时间点"""
        for crop in self.crops:
            # 计算各生长阶段的开始日期
            sowing_start = self.plant_date
            flowering_start = sowing_start + timedelta(days=self.crops[crop]['sowing_days'])
            mature_start = flowering_start + timedelta(days=self.crops[crop]['flowering_days'])
            harvest_start = mature_start + timedelta(days=self.crops[crop]['mature_days'])

            # 存储生长周期
            self.crops[crop]['growth_period'] = {
                '播种期': (sowing_start, flowering_start),
                '开花期': (flowering_start, mature_start),
                '成熟期': (mature_start, harvest_start)
            }

            # 打印生长周期信息
            print(f"{crop}生长周期：")
            for stage, (start, end) in self.crops[crop]['growth_period'].items():
                print(f"  {stage}: {start.strftime('%Y-%m-%d')} 至 {end.strftime('%Y-%m-%d')}")

    def generate_soil_data(self):
        """生成模拟土壤湿度数据（实际应从附件读取）"""
        dates = pd.date_range(start='2021-05-01', end='2021-07-31')
        n_days = len(dates)

        # 模拟土壤湿度（在0.2-0.4之间波动）
        np.random.seed(42)
        base_moisture = 0.3
        daily_fluctuations = np.random.normal(0, 0.03, n_days)
        soil_moisture = np.clip(base_moisture + daily_fluctuations, 0.2, 0.4)

        return pd.DataFrame({
            '日期': dates,
            '土壤湿度': soil_moisture
        })

    def generate_rainfall_data(self):
        """生成模拟降水量数据（实际应从附件读取）"""
        dates = pd.date_range(start='2021-05-01', end='2021-07-31')
        n_days = len(dates)

        # 模拟降水量（大部分天为0，偶尔有降雨）
        np.random.seed(42)
        rainfall = np.zeros(n_days)
        # 随机选择20天有降雨，降雨量在5-30mm之间
        rainy_days = np.random.choice(n_days, 20, replace=False)
        rainfall[rainy_days] = np.random.uniform(5, 30, len(rainy_days))

        return pd.DataFrame({
            '日期': dates,
            '降水量(mm)': rainfall
        })

    def calculate_daily_water_demand(self, date):
        """计算指定日期的总灌溉需水量"""
        total_demand = 0

        # 中文阶段到英文键名的映射
        stage_mapping = {
            '播种期': 'sowing',
            '开花期': 'flowering',
            '成熟期': 'mature'
        }

        for crop, info in self.crops.items():
            # 检查该日期属于哪个生长阶段
            current_stage = None
            for stage, (start, end) in info['growth_period'].items():
                if start <= date < end:
                    current_stage = stage
                    break

            if current_stage:
                # 根据生长阶段和面积计算需水量
                english_stage = stage_mapping.get(current_stage, '')
                if english_stage:
                    water_rate = info[f"{english_stage}_water"]
                    crop_demand = water_rate * info['area'] * 10000  # 转换为平方米
                    total_demand += crop_demand

        return total_demand

    def calculate_monthly_irrigation(self):
        """计算每月灌溉方案"""
        # 划分月份
        months = [
            ('2021-05-01', '2021-05-31'),
            ('2021-06-01', '2021-06-30'),
            ('2021-07-01', '2021-07-31')
        ]

        # 可用储水量（扣除应急储备）
        available_tank_volume = self.system_config['tank_volume'] * (1 - self.system_config['emergency_ratio'])
        current_tank_level = available_tank_volume  # 初始储水量

        for month_start, month_end in months:
            start_date = datetime.strptime(month_start, '%Y-%m-%d')
            end_date = datetime.strptime(month_end, '%Y-%m-%d')
            days_in_month = (end_date - start_date).days + 1

            # 月度统计
            total_irrigation = 0
            river_water = 0
            tank_water = 0
            rainfall_contribution = 0

            # 遍历月份中的每一天
            for day_offset in range(days_in_month):
                current_date = start_date + timedelta(days=day_offset)

                # 1. 计算当日需水量
                daily_demand = self.calculate_daily_water_demand(current_date)

                # 2. 检查降水量（1mm降水量=1L/m²）
                rainfall = self.rainfall_data[
                    self.rainfall_data['日期'] == current_date
                    ]['降水量(mm)'].values[0]

                # 降水量转化为灌溉量等效值
                rainfall_equivalent = rainfall * 10000 * 1.0  # 10000m²（1公顷）的水量
                rainfall_contribution += rainfall_equivalent

                # 3. 实际需灌溉量 = 需水量 - 降水量（假设降水量可覆盖全部农场）
                actual_demand = max(0, daily_demand - rainfall_equivalent)

                # 4. 水源分配（优先使用河水，不足时用储水罐）
                river_supply = min(actual_demand, self.system_config['pipe_daily_flow'])
                tank_supply = max(0, actual_demand - river_supply)

                # 检查储水罐是否足够
                if tank_supply > current_tank_level:
                    # 储水罐不足，调整分配
                    tank_supply = current_tank_level
                    river_supply = actual_demand - tank_supply

                # 更新储水罐水位
                current_tank_level -= tank_supply

                # 累加统计
                total_irrigation += actual_demand
                river_water += river_supply
                tank_water += tank_supply

            # 计算水源比例
            if total_irrigation > 0:
                river_percentage = river_water / total_irrigation * 100
                tank_percentage = tank_water / total_irrigation * 100
                rainfall_percentage = rainfall_contribution / (total_irrigation + rainfall_contribution) * 100
            else:
                river_percentage = 0
                tank_percentage = 0
                rainfall_percentage = 0

            # 检查系统是否满足需求
            system_adaptable = current_tank_level >= 0

            # 存储月度结果
            self.monthly_results.append({
                '月份': f"{start_date.month}月",
                '总灌溉量(L)': total_irrigation,
                '河水比例(%)': river_percentage,
                '储水罐比例(%)': tank_percentage,
                '降水量贡献(%)': rainfall_percentage,
                '系统是否满足需求': system_adaptable,
                '月末储水罐水位(L)': current_tank_level
            })

            # 如果系统不满足需求，调整配置（这里简化为增加管道流量）
            if not system_adaptable:
                self.system_config['pipe_daily_flow'] *= 1.2
                print(
                    f"⚠️ {start_date.month}月系统供水不足，已将管道流量增加20%至{self.system_config['pipe_daily_flow']:.0f}L/日")

        return pd.DataFrame(self.monthly_results)

    def visualize_monthly_irrigation(self):
        """可视化月度灌溉方案"""
        if not self.monthly_results:
            print("请先计算月度灌溉方案")
            return

        # 提取数据
        months = [r['月份'] for r in self.monthly_results]
        total_irrigation = [r['总灌溉量(L)'] for r in self.monthly_results]
        river_percentage = [r['河水比例(%)'] for r in self.monthly_results]
        tank_percentage = [r['储水罐比例(%)'] for r in self.monthly_results]
        rainfall_percentage = [r['降水量贡献(%)'] for r in self.monthly_results]

        # 创建图表
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 10))

        # 1. 月度总灌溉量柱状图
        bars = ax1.bar(months, total_irrigation, color='skyblue')
        ax1.set_title('月度总灌溉量')
        ax1.set_ylabel('灌溉量 (L)')
        ax1.bar_label(bars, padding=3, fmt='%.0f')

        # 2. 水源比例堆叠柱状图
        bottom = np.zeros(len(months))
        ax2.bar(months, river_percentage, label='河水', color='blue', bottom=bottom)
        bottom += np.array(river_percentage)
        ax2.bar(months, tank_percentage, label='储水罐', color='green', bottom=bottom)
        bottom += np.array(tank_percentage)
        ax2.bar(months, rainfall_percentage, label='降水量贡献', color='lightblue', bottom=bottom)
        ax2.set_title('灌溉水源比例')
        ax2.set_ylabel('比例 (%)')
        ax2.legend(loc='upper right')

        plt.tight_layout()
        plt.show()

    def visualize_crop_growth(self):
        """可视化作物生长周期与灌溉需求"""
        # 创建图表
        fig, ax = plt.subplots(figsize=(12, 6))

        # 定义颜色映射
        colors = {
            '高粱': '#FF7F50',  # 橙色
            '玉米': '#FFD700',  # 金色
            '大豆': '#228B22'  # 绿色
        }

        # 绘制每个作物的生长阶段
        for crop, info in self.crops.items():
            y_pos = list(self.crops.keys()).index(crop)

            for stage, (start, end) in info['growth_period'].items():
                # 计算阶段持续天数
                days = (end - start).days

                # 设置阶段颜色
                if stage == '播种期':
                    stage_color = colors[crop]
                elif stage == '开花期':
                    stage_color = colors[crop]
                else:  # 成熟期
                    stage_color = colors[crop]

                # 绘制阶段条
                ax.barh(y_pos, days, left=mdates.date2num(start),
                        height=0.5, color=stage_color, alpha=0.8,
                        label=f"{crop}-{stage}" if stage == '播种期' else "")

                # 添加阶段标签
                ax.text(mdates.date2num(start) + days / 2, y_pos, stage,
                        ha='center', va='center', color='black', fontweight='bold')

        # 设置Y轴标签
        ax.set_yticks(range(len(self.crops)))
        ax.set_yticklabels(self.crops.keys())

        # 设置X轴为日期格式
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        ax.xaxis.set_major_locator(mdates.DayLocator(interval=10))
        plt.xticks(rotation=45)

        # 添加标题和图例
        ax.set_title('作物生长周期与灌溉阶段')
        ax.legend(loc='upper right')

        plt.tight_layout()
        plt.show()

    def generate_table5(self):
        """生成表5结果"""
        if not self.monthly_results:
            print("请先计算月度灌溉方案")
            return None

        # 创建表格数据
        table5 = pd.DataFrame(self.monthly_results)

        # 添加系统调整说明
        table5['系统调整说明'] = [
            "无需调整" if row['系统是否满足需求'] else f"增加管道流量至{self.system_config['pipe_daily_flow']:.0f}L/日"
            for _, row in table5.iterrows()
        ]

        # 调整列顺序
        table5 = table5[['月份', '总灌溉量(L)', '河水比例(%)', '储水罐比例(%)', '降水量贡献(%)', '系统是否满足需求',
                         '系统调整说明']]

        return table5

    def run_analysis(self):
        """执行问题4完整分析流程"""
        print("=== 问题4：作物月度灌溉方案规划 ===")

        # 步骤1：计算月度灌溉方案
        monthly_results = self.calculate_monthly_irrigation()
        print("\n月度灌溉方案计算结果：")
        print(monthly_results)

        # 步骤2：生成表5
        table5 = self.generate_table5()
        print("\n表5 月度灌溉方案表：")
        print(table5)

        # 步骤3：可视化月度灌溉情况
        self.visualize_monthly_irrigation()

        # 步骤4：可视化作物生长周期
        self.visualize_crop_growth()

        return {
            '月度结果': monthly_results,
            '表5': table5
        }


if __name__ == "__main__":
    # 模拟问题2的结果
    problem2_result = {
        '管道日流量': 5000,  # 管道日流量（L）
        '储水罐容积': 30000,  # 储水罐总容积（L）
        '管道长度': 80,  # 管道长度（米）
        '喷头覆盖半径': 15  # 喷头覆盖半径（米）
    }

    # 模拟问题3的结果
    problem3_result = {
        'emergency_ratio': 0.3  # 应急储备比例
    }

    # 创建灌溉规划器并运行分析
    planner = CropIrrigationPlanner(problem2_result, problem3_result)
    result = planner.run_analysis()